from data_process.path_algorithm.PathAlgorithm import PathAlgorithm
from constants import *
import os
import pickle
from utils.file_util import *

class MultiDirSameAlgorithm(PathAlgorithm):
    def __init__(self, params):
        self.params = params
        self.rootPath = params.rootPath
        self.volun = params.volun
        self.trainDirList = params.trainDirList
        self.testDirList = params.testDirList
        self.vocabPath = params.vocabPath
        self.mapPath = params.mapPath
        self.trainPath = TRAIN_PATH
        self.testPath = TEST_PATH

    def generateTimeMarkFile(self):
        pass

    def generatePath(self):
        self.splitPathFromAllDir()
        mapTabel = CreateMapTable(self.mapPath)
        createContentFromFilePath(self.trainPath, self.testPath, mapTabel)
        self.generateTimeMarkFile()

    def splitPathFromAllDir(self):
        trainData = []
        testData = []
        for dirPath in self.trainDirList:
            SplitDataPath(dirPath, self.vocabPath, 0.2)
            self.selectDataFromFile(self.trainPath)
            self.addAllPrefix(dirPath)
            tempTrain = readDataFromFile(self.trainPath)
            tempTest = readDataFromFile(self.testPath)
            trainData += tempTrain
            testData += tempTest
        with open(self.trainPath, 'w') as f:
            for path in trainData:
                f.write(f'{path}\n')
        with open(self.testPath, 'w') as f:
            for path in testData:
                f.write(f'{path}\n')

    def addAllPrefix(self, prefix):
        self.addPrefixOfPathFile(self.trainPath, prefix)
        self.addPrefixOfPathFile(self.testPath, prefix)

    def addPrefixOfPathFile(self, filePath, prefix):
        savePaths = []
        paths = readDataFromFile(filePath)
        for curPath in paths:
            res = os.path.join(prefix, curPath)
            savePaths.append(res)
        with open(filePath, 'w') as f:
            for curPath in savePaths:
                f.write(f'{curPath}\n')

    def selectDataFromFile(self, filePath, selectRatio=1.0):
        allPath = readDataFromFile(filePath)
        prefixDict = {}
        trainPath = []
        selectNumber = int(selectRatio * len(allPath))
        for curPath in allPath:
            pre = self.getPrefixFromPath(curPath)
            if pre not in prefixDict.keys():
                prefixDict[pre] = []
            prefixDict[pre].append(curPath)
        for curPre in prefixDict.keys():
            curList = prefixDict[curPre]
            curTrain = []
            if selectNumber>0:
                fNum = min(len(curList), selectNumber)
                offList = random.sample(range(0, len(curList)), fNum)
                for i in offList:
                    curTrain.append(curList[i])
            else:
                curTrain = curList
            trainPath += curTrain

        with open(filePath, 'w') as f:
            for path in trainPath:
                f.write(f'{path}\n')

    def getPrefixFromPath(self, path):
        pathLen = len(path)
        prefix = ''
        for i in range(pathLen-1, 0, -1):
            if(path[i] == '_'):
                prefix = path[0:i]
                break
        return prefix